import requests
import json
from typing import Union, Dict

API_KEY = "UHRPfgJhWiFIv4DNCR13Q90z"
SECRET_KEY = "oi04isr4nITGRp1GjIyPzjcclmcD6W0K"


def send_req(msg: Union[str, Dict]):
    if isinstance(msg, str):
        msg = {"messages": [{"role": "user", "content": msg}]}
    url = (
        "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions?access_token="
        + get_access_token()
    )

    payload = {
        "temperature": 0.95,
        "top_p": 0.8,
        "penalty_score": 1,
        "disable_search": False,
        "enable_citation": False,
        "response_format": "text",
    }
    payload.update(msg)
    payload = json.dumps(payload)
    headers = {"Content-Type": "application/json"}

    response = requests.request("POST", url, headers=headers, data=payload)

    # return response.text
    return json.loads(response.text)


def get_access_token():
    """
    使用 AK，SK 生成鉴权签名（Access Token）
    :return: access_token，或是None(如果错误)
    """
    url = "https://aip.baidubce.com/oauth/2.0/token"
    params = {
        "grant_type": "client_credentials",
        "client_id": API_KEY,
        "client_secret": SECRET_KEY,
    }
    return str(requests.post(url, params=params).json().get("access_token"))


if __name__ == "__main__":
    send_req()
